import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader

from pokemon import PokemonData
from model import Digit

BATCH_SIZE = 10
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCHS = 2

train_set = PokemonData('train', 100, 'train')
test_set = PokemonData('test', 100, 'test')

train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)

# s
model = Digit().to(DEVICE)
optimizer = optim.Adam(model.parameters())


def train_model(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_index, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        # _,pred = output.max(1, keepdim=True)
        # or pred=output.max(1)[1]
        loss.backward()
        optimizer.step()
        if batch_index % 10 == 0:
            print("Train epoch : {} \t loss : {:.6f}".format(
                epoch, loss.item()))
        # print(data.shape)
        # model.save('./model.pt')
        # model.save('./optimizer.pt')


def test_model(model, device, test_loader):
    model.eval()
    correct = 0.0
    test_loss = 0.0
    with torch.no_grad():  # 不计算梯度，不会反向传播
        for batch_index, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target).item()
            pred = output.max(1, keepdim=True)[1]  # [1] is pred.indice
            correct += pred.eq(target.view_as(pred)).sum().item()
            print(pred.T)
            print(target)
        test_loss /= len(test_loader.dataset)
        print("Test --ave loss : {: .4f} , Accuracy : {:.3f}%\n".format(
            test_loss, 100*correct/len(test_loader.dataset)
        ))


example = torch.rand(10, 1, 100, 100)

for epoch in range(1, EPOCHS+1):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_model(model, DEVICE, test_loader)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("./model.pt")
